package cn.jinbyte.web.wrapper;


import jakarta.servlet.ReadListener;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.ServletRequest;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import lombok.Getter;
import org.jetbrains.annotations.NotNull;
import org.springframework.util.FastByteArrayOutputStream;
import org.springframework.util.StreamUtils;

import java.io.*;
import java.net.URLEncoder;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

/**
 * 重复读取请求内容
 *
 * @author JinByte
 */

public class RepeatReadRequestWrapper extends HttpServletRequestWrapper {
    @Getter
    private final FastByteArrayOutputStream cachedContent;
    private ContentCachingInputStream contentCachingInputStream;
    private final String characterEncoding;

    public RepeatReadRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        this.characterEncoding = request.getCharacterEncoding() != null ? request.getCharacterEncoding() : StandardCharsets.UTF_8.name();
        int contentLength = super.getRequest().getContentLength();
        this.cachedContent = contentLength > 0 ? new FastByteArrayOutputStream(contentLength) : new FastByteArrayOutputStream();
        if (!this.isMultipartContent(request)) {
            if (this.isFormRequest()) {
                this.writeRequestParametersToCachedContent();
            } else {
                StreamUtils.copy(request.getInputStream(), this.cachedContent);
            }

            this.contentCachingInputStream = new ContentCachingInputStream(this.cachedContent.toByteArray());
        }

    }

    public ServletInputStream getInputStream() throws IOException {
        if (this.isMultipartContent(super.getRequest())) {
            return super.getRequest().getInputStream();
        } else {
            synchronized (this) {
                this.contentCachingInputStream.reset();
                return this.contentCachingInputStream;
            }
        }
    }

    public BufferedReader getReader() throws IOException {
        if (this.isMultipartContent(super.getRequest())) {
            return super.getRequest().getReader();
        } else {
            synchronized (this) {
                return new BufferedReader(new InputStreamReader(this.getInputStream(), this.getCharacterEncoding()));
            }
        }
    }

    private void writeRequestParametersToCachedContent() {
        try {
            if (this.cachedContent.size() == 0) {
                Map<String, String[]> form = super.getParameterMap();
                Iterator<String> nameIterator = form.keySet().iterator();

                while (nameIterator.hasNext()) {
                    String name = (String) nameIterator.next();
                    List<String> values = Arrays.asList((String[]) form.get(name));
                    Iterator<String> valueIterator = values.iterator();

                    while (valueIterator.hasNext()) {
                        String value = (String) valueIterator.next();
                        this.cachedContent.write(URLEncoder.encode(name, this.characterEncoding).getBytes());
                        if (value != null) {
                            this.cachedContent.write("=".getBytes());
                            this.cachedContent.write(URLEncoder.encode(value, this.characterEncoding).getBytes());
                            if (valueIterator.hasNext()) {
                                this.cachedContent.write("&".getBytes());
                            }
                        }
                    }

                    if (nameIterator.hasNext()) {
                        this.cachedContent.write("&".getBytes());
                    }
                }
            }

        } catch (IOException ex) {
            throw new IllegalStateException("Failed to write request parameters to cached content", ex);
        }
    }

    @Override
    public String getCharacterEncoding() {
        return this.characterEncoding;
    }

    public String getContentAsString() {
        return this.cachedContent.toString(Charset.forName(this.getCharacterEncoding()));
    }

    public boolean isMultipartContent(ServletRequest request) {
        String contentType = request.getContentType();
        return contentType != null && contentType.toLowerCase().startsWith("multipart/");
    }

    private boolean isFormRequest() {
        String contentType = this.getContentType();
        return contentType != null && contentType.contains("application/x-www-form-urlencoded");
    }

    private static class ContentCachingInputStream extends ServletInputStream {
        private final InputStream delegate;

        public ContentCachingInputStream(byte[] body) {
            this.delegate = new ByteArrayInputStream(body);
        }

        public boolean isFinished() {
            return false;
        }

        public boolean isReady() {
            return true;
        }

        public void setReadListener(ReadListener readListener) {
            throw new UnsupportedOperationException();
        }

        public int read() throws IOException {
            return this.delegate.read();
        }

        public int read(byte @NotNull [] b, int off, int len) throws IOException {
            return this.delegate.read(b, off, len);
        }

        public int read(byte @NotNull [] b) throws IOException {
            return this.delegate.read(b);
        }

        public long skip(long n) throws IOException {
            return this.delegate.skip(n);
        }

        public int available() throws IOException {
            return this.delegate.available();
        }

        public void close() throws IOException {
            this.delegate.close();
        }

        public synchronized void mark(int readlimit) {
            this.delegate.mark(readlimit);
        }

        public synchronized void reset() throws IOException {
            this.delegate.reset();
        }

        public boolean markSupported() {
            return this.delegate.markSupported();
        }
    }
}
